import my_reduce
import torch


class MyReduceFunc(torch.autograd.Function):
    @staticmethod
    def forward(ctx, y):
        ctx.save_for_backward(y)
        return my_reduce.forward(y)  # 调 C++ 实现

    @staticmethod
    def backward(ctx, grad_out):
        (y,) = ctx.saved_tensors
        return my_reduce.backward(grad_out, y)


class MyReduce(torch.nn.Module):
    def forward(self, y):
        return MyReduceFunc.apply(y)
